#include <iostream>
using namespace std;
#include <stdio.h>
#include <cmath>
#include <mex.h>

#include <vector>
#include <algorithm>
#include <queue>
#include <limits>
#include <mexFileIO.h>
#include <szMexUtility.h>
#include <szMexUtilityTemplate.h>
#include <szMyNeighborOp.h>
#include <szMiscOperations.h>
#include <szParticleF.h>
#include <FragmentInfo.h>
#include <DistanceMapUtility.h>
#include <Graph.h>
#include <GraphFactory.h>
#include <Kruskal.h>

GraphFactory<CParticleF>* GraphFactory<CParticleF>::_instance = NULL;

void
doMedialAxis(int nlhs, mxArray **plhs, int nrhs, const mxArray *prhs[])
{
	printf("Doing Medial Axis Transform.\n");
	if(nrhs < 2)
	{
		mexErrMsgTxt("Usage: [P mst] = DotOperations('mat', dmap)");
		return;
	}
	vector<float> dmap;
	const int* dims;
	{
		mxClassID classIdP;
		int ndimP;
		LoadData(dmap, prhs[1], classIdP, ndimP, &dims);
	}

	vector<float> cent(dmap.size(), 0.0f);
	CentricityTransform(cent, dmap, dims);
	vector<CParticleF> ridges;
	for(int y=0; y<dims[1]; ++y)
	{
		for(int x=0; x<dims[0]; ++x)
		{
			float c = GetData2(cent, x, y, dims[0], dims[1], 0.0f);
			if(c >= 6.0/8.0)
			{
				ridges.push_back(CParticleF(x, y, 0, c));
			}
		}
	}

	vector<Vertex<CParticleF>*> vertices(ridges.size());
	GraphFactory<CParticleF>* factory = GraphFactory<CParticleF>::GetInstance();
	for(int i=0; i<ridges.size(); ++i)
	{
		vertices[i] = factory->makeVertex(ridges[i]);
	}
	vector<Edge<CParticleF>*> edges;
	for(int i=0; i<ridges.size(); ++i)
	{
		for(int j=0; j<ridges.size(); ++j)
		{
			float d = Distance(ridges[i], ridges[j]);
			Edge<CParticleF>* edge = factory->makeEdge(vertices[i], vertices[j], d*d);
			vertices[i]->Add(edge);
			edges.push_back(edge);
		}
	}

	vector<Edge<CParticleF>*> mst = Kruskal(edges, vertices);
	if(nlhs >= 1)
	{
		//write x, y, centricity value, distance value
		const int dims[] = {ridges.size(), 4};
		vector<float> F(dims[0]*dims[1]);
		for(int i=0; i<dims[0]; ++i)
		{
			SetData2(F, i, 0, dims[0], dims[1], ridges[i].m_X);
			SetData2(F, i, 1, dims[0], dims[1], ridges[i].m_Y);
			SetData2(F, i, 2, dims[0], dims[1], ridges[i].m_Life);
			float dval = GetData2(dmap, ridges[i].m_X, ridges[i].m_Y, dims[0], dims[1], 0.0f);
			SetData2(F, i, 3, dims[0], dims[1], dval);
		}
		plhs[0] = StoreData(F, mxSINGLE_CLASS, 2, dims);
	}
	if(nlhs >= 2)
	{
		//indices to two ridge points, plus the weight
		const int dims[] = {mst.size(), 3};
		vector<float> F(dims[0]*dims[1]);
		for(int i=0; i<dims[0]; ++i)
		{
			int k1 = distance(ridges.begin(), find(ridges.begin(), ridges.end(), mst[i]->u->key));
			int k2 = distance(ridges.begin(), find(ridges.begin(), ridges.end(), mst[i]->v->key));
			SetData2(F, i, 0, dims[0], dims[1], (float)k1+1);
			SetData2(F, i, 1, dims[0], dims[1], (float)k2+1);
			SetData2(F, i, 2, dims[0], dims[1], mst[i]->w);
		}
		plhs[1] = StoreData(F, mxSINGLE_CLASS, 2, dims);
	}

	factory->Clean();
}

